package com.ics.atable.chat.node;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.ics.atable.chat.config.ChatClientServiceFactory;
import com.ics.atable.chat.model.dto.IntentExtract;
import com.ics.atable.chat.model.dto.KeyWordExtract;
import com.ics.atable.chat.prompt.PromptConstant;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.ics.atable.chat.constant.ItemQueryConstant.*;

import java.util.HashMap;
import java.util.Map;

@Slf4j
public class StructTokenNode implements NodeAction {

    private final ChatClient structTokenClient;

    public StructTokenNode(ChatClientServiceFactory chatClientServiceFactory) {
        this.structTokenClient = chatClientServiceFactory.getOrCreateStructTokenChatClient();
    }


    @Override
    public Map<String, Object> apply(OverAllState state) throws Exception {
        Map<String, Object> updated = new HashMap<>();
        try{

            IntentExtract intentExtract = state.value("intent_result", IntentExtract.class)
                    .orElseThrow(() -> new IllegalArgumentException("intent is missing"));
            log.info("intentExtract：{}", intentExtract);
            String prompt = intentExtract.getOutput();

            // 拿到其中为String的Template
            String systemPrompt = PromptConstant.getStructTokenPromptTemplate().getTemplate();
            KeyWordExtract llm_response = structTokenClient.prompt().system(systemPrompt).user(prompt).call().entity(KeyWordExtract.class);
            log.info("llm_response: {}", llm_response);

            updated.put("struct_token", llm_response);
            updated.put(STRUCT_TOKEN_NEXT_NODE, SYNONYMS);
            // 连接下一个并行检索结点 (Parallel search)
//            updated.put("parallel_search", jsonString);
            return updated;
        } catch (Exception e){
            log.error("struct_token node error: {}", e.getMessage(), e);
            updated.put("struct_token", "");
            // 出现错误直接结束
            updated.put(STRUCT_TOKEN_NEXT_NODE, END);
            updated.put("error", e.getMessage());
            return updated;
        }

    }
}
